# data loader
from __future__ import print_function, division
import os
import torch
from skimage import io, transform, color
import numpy as np
from torch.utils.data import Dataset, DataLoader
from osgeo import gdal

#==========================dataset load==========================
class ToTensorNorm(object):
    """Convert ndarrays in sample to Tensors."""
    def __init__(self):
        pass
    def __call__(self, sample):
        
        image, label = sample['image'], sample['mask']
        
        if (np.max(label) < 1e-6):
            label = label
        else:
            label = label / np.max(label)
        
        if np.max(image) != 0:
            image = image / np.max(image)
        else:
            image = image

        return {'image': torch.from_numpy(image),
                'mask': torch.from_numpy(label)}

class CloudDataset(Dataset):
    def __init__(self,file_name_list,transform=None,test_mode=0):
        self.file_name_list = file_name_list
        # self.label_name_list = lbl_nam
        # e_list
        self.transform = transform
        self.root = '/home/data/wukang/datasets/l8/'
        self.test_mode = test_mode

    def __len__(self):
        return len(self.file_name_list)

    def __getitem__(self,idx):
        if self.test_mode ==0:
            img_name, label_name = self.file_name_list[idx].split(' ')
            label_name = label_name.replace("\n", "")
        else:
            img_name = self.file_name_list[idx]
            label_name = None
        image_dataset = gdal.Open(os.path.join(self.root, img_name))
        width = image_dataset.RasterXSize
        height = image_dataset.RasterYSize
        img_10ch = image_dataset.ReadAsArray(0, 0, width, height)
        #label_1ch = cv2.imread(os.path.join(self.root, label_name),cv2.IMREAD_GRAYSCALE)
        if self.test_mode == 0:
            label_1ch = io.imread(os.path.join(self.root, label_name), as_gray=1)
        else:
            label_1ch = np.zeros((width,height))
        label_1ch = label_1ch[np.newaxis, :, :]

        img_10ch = img_10ch.astype(np.float32)
        label_1ch = label_1ch.astype(np.float32)

        sample = {'image':img_10ch, 'mask':label_1ch}

        if self.transform:
            sample = self.transform(sample)

        return sample